#include <stdlib.h>
#include "TU.h"
#include "CFG.h"

TranslationUnit* TU_New(){
    TranslationUnit* New;
    
    New = malloc(sizeof(TranslationUnit));
    
    New->AddressFunc = NULL;
    New->CFGRepoRefreshFunc = NULL;
    New->Type = TUT_NONE;
    return New;
}

static uint16_t TU_BBAddress(TranslationUnit* TU){
    return TU->NodeData.BB->StartingAddress;
}

static uint16_t TU_CFGAddress(TranslationUnit *TU){
    return TU->NodeData.CFG->InitialAddress;
}


static void TU_CFGRBBRefresh(TranslationUnit *TU, ControlFlowGraph *CFG){
    CFGRBBRefresh(CFG, TU->NodeData.BB);
}

static void TU_CFGRCFGRefresh(TranslationUnit *TU, ControlFlowGraph *CFG){
    CFGRRefresh(TU->NodeData.CFG);
}

static uint64_t TU_BBSize(TranslationUnit* TU){
    return TU->NodeData.BB->NumInstructions;
}

static uint64_t TU_CFGSize(TranslationUnit* TU){
    ControlFlowGraph* CFG;
    CFGNode* Node;
    uint64_t Size = 0;
    
    CFG = TU->NodeData.CFG;
    
    LST_Rewind(CFG->NodeList);
    while((Node = LST_Next(CFG->NodeList))){
        Node->_VisitedSize = 0;
    }
    
    LST_Rewind(CFG->NodeList);
    
    while((Node = LST_Next(CFG->NodeList))){
        if(!Node->_VisitedSize)
            Node->_VisitedSize = 1,
                Size += TU_GetSize(Node->Unit);
        else
            Size += Node->Size;
    }
    
    return Size;
}

void TU_CFGRefresh(TranslationUnit* TU, ControlFlowGraph* CFG){
    TU->CFGRepoRefreshFunc(TU, CFG);
}

void TU_SetBB(TranslationUnit* TU, BasicBlock* BB){
    TU->NodeData.BB = BB;
    TU->Type = TUT_BB;
    TU->AddressFunc = TU_BBAddress;
    TU->CFGRepoRefreshFunc = TU_CFGRBBRefresh;
    TU->SizeFunc = TU_BBSize;
}

void TU_SetCFG(TranslationUnit* TU, ControlFlowGraph* CFG){
    TU->NodeData.CFG = CFG;
    TU->Type = TUT_CFG;
    TU->AddressFunc = TU_CFGAddress;
    TU->CFGRepoRefreshFunc = TU_CFGRCFGRefresh;
    TU->SizeFunc = TU_CFGSize;
}

BasicBlock* TU_GetBB(TranslationUnit* TU){
    if(TU->Type == TUT_BB)
        return TU->NodeData.BB;
    else
        return NULL;
}

ControlFlowGraph* TU_GetCFG(TranslationUnit* TU){
    if(TU->Type == TUT_CFG)
        return TU->NodeData.CFG;
    else
        return NULL;
}

uint16_t TU_GetStartingAddress(TranslationUnit* TU){
    return TU->AddressFunc(TU);
}

uint32_t TU_IsSameTU(TranslationUnit* TU1, TranslationUnit* TU2){
    if(TU1->Type == TU2->Type)
        if(TU1->AddressFunc(TU1) == TU2->AddressFunc(TU2))
            return 1;
    return 0;
}

uint32_t TU_IsNone(TranslationUnit* TU){
    return TU->Type == TUT_NONE ? 1 : 0;
}

uint32_t TU_GetSize(TranslationUnit* TU){
    return TU->SizeFunc(TU);
}
